'''
羊吃草教程。注意导入方式和之前几个示例的区别——可以自由灵活选用各种导入方式，效果都是相同的。
'''
from takagiabm import GridAgent, simStart, GridModel

import numpy as np
from takagiabm.agents.gridagents.freegrid.freegridcell import Cell
from takagiabm.variable import Var
from takagiabm.datacounter import DataCounter
import takagiabm as tak


def createSheep():  # 调用这个函数生成一只“羊”
    a = GridAgent()
    triangle2 = tak.scale(tak.takArrow, 1)  # 设置形状为takagi库的箭头形状，缩放比例为3.
    a.setTrait('energy', 10)
    a.setShape(triangle2)  # 设置形状为三角形。
    a.setSpeed(np.array([1, 0]))  # 设置初始速度
    a.setColor(tak.red)  # 设置颜色
    return a


def modelSetup(model: GridModel):# 如果引入类型声明，则可以让IDE自动判断出变量的类型，从而辅助自动补全。
    # 由此，我推荐使用类型声明。
    for i in range(5):
        a = createSheep()
        model.addAgent(a, randomPlace=True)

    for i in range(model.width):
        for j in range(model.height):
            if (np.random.random() < 0.5):
                model.grid.cells[j][i].setColor(tak.green) # color是保留属性，不允许写成setTrait('color',tak.red)。


def agentStep(agent):  # 代理人(也就是羊)执行的函数
    global stat, sheepGrowth
    color = agent.getCellColor()
    if (color == tak.green):
        agent.setCellColor(tak.white)
        agent.properties['energy'] += 5

    if (np.random.random() < sheepGrowth.value):
        if (agent.getTrait('energy') > 10):
            a = createSheep()
            a.setTrait('energy', 5)
            agent.setTrait('energy', agent.getTrait('energy') - 5)
            a.setPos(agent.pos.astype(np.float))
            agent.model.addAgent(a)

    energy = agent.getTrait('energy')
    agent.setTrait('energy', energy - 1)  # 这种是get和set的风格。也可以直接写成agent.traits['energy']-=1
    if (agent.getTrait('energy') == 0):
        agent.model.removeAgent(agent)
        return
    agent.turn(np.random.random() * np.pi - np.pi / 2)
    agent.move()


def cellStep(cell):
    global grassGrowth
    if (np.random.random() < grassGrowth.value):
        if (cell.getColor() == tak.white):
            cell.setColor(tak.green)


def cellOnClicked(cell):  # 当网格被点击的时候，就变换一下颜色——产生一些草。
    if (cell.getColor() == tak.white):
        cell.setColor(tak.green)
    else:
        cell.setColor(tak.white)


sheepGrowth = Var('sheepBirth', 0.03, range=(0, 0.01, 0.05))
grassGrowth = Var('grassGrowth', 0.05, range=(0, 0.01, 0.5))

varList = [grassGrowth, sheepGrowth]
counterList = [DataCounter(name='颜色', propertyName='color', targets=['cell']),
               DataCounter(propertyName='color', targets=['agent'])]

Cell = tak.prepareCell(Cell, cellStepFunc=cellStep, cellOnClicked=cellOnClicked)
GridAgent = tak.prepareAgent(GridAgent, agentStepFunc=agentStep)
GridModel = tak.prepareModel(GridModel, varList=varList, dataCounterList=counterList, width=34, height=34,
                             modelInitFunc=modelSetup, cellActivationPolicy='casual', agentActivationPolicy='casual'
                             )
if __name__ == "__main__":
    simStart(__file__, GridModel, maxSteps=10000)  # 这个方法同时启动了图形界面。
